#include "xict_cpp/core/stage/stream_processor.h"

namespace xict_cpp {
    namespace core {
        namespace stage {
            bool StreamProcessorStage::Init() {
                StreamProcessorStageParams params;
                return Init(params);
            }

            bool StreamProcessorStage::Init(
                const StreamProcessorStageParams& params) {
                params_ = params;
                return Initialize();
            }

            bool StreamProcessorStage::Initialize() {
                if (params_.model_flag) {
                    detector_ptr_ = std::make_shared<
                        xict_common::perception::detection::Yolov5TrtDet>(
                        params_.model_path);

                    if (params_.ignore_region_flag) {
                        // // 创建一个掩码，初始化为全黑
                        // mask_ = cv::Mat::zeros(params_.image_height,
                        //                        params_.image_width, CV_8UC1);
                        // // 填充忽略区域
                        // for (const auto& region : params_.ignore_regions) {
                        //     // 用白色填充多边形区域
                        //     cv::fillPoly(
                        //         mask_,
                        //         std::vector<std::vector<cv::Point>> {region},
                        //         cv::Scalar(255));
                        // }
                        // if (params_.mask_show_flag) {
                        //     cv::imshow("mask", mask_);
                        //     cv::waitKey(0);
                        // }
                    }
                }

                image_subscriber_.subscriber_front_image = std::make_shared<
                    eCAL::protobuf::CSubscriber<xcmg_proto::OpencvImage>>(
                    params_.front_image_topic);
                image_subscriber_.subscriber_back_image = std::make_shared<
                    eCAL::protobuf::CSubscriber<xcmg_proto::OpencvImage>>(
                    params_.back_image_topic);
                image_subscriber_.subscriber_left_image = std::make_shared<
                    eCAL::protobuf::CSubscriber<xcmg_proto::OpencvImage>>(
                    params_.left_image_topic);
                image_subscriber_.subscriber_right_image = std::make_shared<
                    eCAL::protobuf::CSubscriber<xcmg_proto::OpencvImage>>(
                    params_.right_image_topic);

                detector_image_publisher_.publisher_yolo_front =
                    std::make_shared<
                        eCAL::protobuf::CPublisher<xcmg_proto::OpencvImage>>(
                        params_.yolo_front);
                detector_image_publisher_.publisher_yolo_back =
                    std::make_shared<
                        eCAL::protobuf::CPublisher<xcmg_proto::OpencvImage>>(
                        params_.yolo_back);
                detector_image_publisher_.publisher_yolo_left =
                    std::make_shared<
                        eCAL::protobuf::CPublisher<xcmg_proto::OpencvImage>>(
                        params_.yolo_left);
                detector_image_publisher_.publisher_yolo_right =
                    std::make_shared<
                        eCAL::protobuf::CPublisher<xcmg_proto::OpencvImage>>(
                        params_.yolo_right);

                trajectory_processor_ = std::make_shared<
                    xict_cpp::core::stage::DrawTrajectoryStage>();

                if (params_.rtsp_enable) {
                    std::string rtsp_source =
                        fmt::format("appsrc ! queue ! videoconvert ! x264enc "
                                    "tune=zerolatency speed-preset=ultrafast");

                    rtsp_publisher_.rtsp_url_front = fmt::format(
                        "{} bitrate={} key-int-max=40 ! rtspclientsink "
                        "location={}/{}",
                        rtsp_source, params_.rtsp_bitrate, params_.rtsp_url,
                        params_.front_image_topic);
                    rtsp_publisher_.rtsp_url_back =
                        fmt::format("{} bitrate={} key-int-max=40 ! "
                                    "rtspclientsink location={}/{}",
                                    rtsp_source, params_.rtsp_bitrate,
                                    params_.rtsp_url, params_.back_image_topic);
                    rtsp_publisher_.rtsp_url_left =
                        fmt::format("{} bitrate={} key-int-max=40 ! "
                                    "rtspclientsink location={}/{}",
                                    rtsp_source, params_.rtsp_bitrate,
                                    params_.rtsp_url, params_.left_image_topic);
                    rtsp_publisher_.rtsp_url_right = fmt::format(
                        "{} bitrate={} key-int-max=40 ! rtspclientsink "
                        "location={}/{}",
                        rtsp_source, params_.rtsp_bitrate, params_.rtsp_url,
                        params_.right_image_topic);

                    rtsp_publisher_.writer_front = cv::VideoWriter(
                        rtsp_publisher_.rtsp_url_front, cv::CAP_GSTREAMER, 0,
                        params_.rtsp_fps,
                        cv::Size(params_.rtsp_width, params_.rtsp_height),
                        true);
                    rtsp_publisher_.writer_back = cv::VideoWriter(
                        rtsp_publisher_.rtsp_url_back, cv::CAP_GSTREAMER, 0,
                        params_.rtsp_fps,
                        cv::Size(params_.rtsp_width, params_.rtsp_height),
                        true);
                    rtsp_publisher_.writer_left = cv::VideoWriter(
                        rtsp_publisher_.rtsp_url_left, cv::CAP_GSTREAMER, 0,
                        params_.rtsp_fps,
                        cv::Size(params_.rtsp_width, params_.rtsp_height),
                        true);
                    rtsp_publisher_.writer_right = cv::VideoWriter(
                        rtsp_publisher_.rtsp_url_right, cv::CAP_GSTREAMER, 0,
                        params_.rtsp_fps,
                        cv::Size(params_.rtsp_width, params_.rtsp_height),
                        true);
                }

                return true;
            }

            std::string StreamProcessorStage::Name() const {
                return "StreamProcessorStage";
            }

            void StreamProcessorStage::Reset() {
                if (detector_ptr_) {
                    detector_ptr_.reset();
                    detector_ptr_ = nullptr;
                }
            }

            bool StreamProcessorStage::Process() {
                auto lambda_front =
                    [&](const char*, const xcmg_proto::OpencvImage& data_front,
                        long long, long long, long long) {
                        CallbackImageDetect(
                            params_.front_image_topic.c_str(), data_front, 0, 0,
                            0, detector_image_publisher_.publisher_yolo_front);
                    };
                image_subscriber_.subscriber_front_image->AddReceiveCallback(
                    lambda_front);

                auto lambda_back = [&](const char*,
                                       const xcmg_proto::OpencvImage& data_back,
                                       long long, long long, long long) {
                    CallbackImageDetect(
                        params_.back_image_topic.c_str(), data_back, 0, 0, 0,
                        detector_image_publisher_.publisher_yolo_back);
                };
                image_subscriber_.subscriber_back_image->AddReceiveCallback(
                    lambda_back);

                auto lambda_left = [&](const char*,
                                       const xcmg_proto::OpencvImage& data_left,
                                       long long, long long, long long) {
                    CallbackImageDetect(
                        params_.left_image_topic.c_str(), data_left, 0, 0, 0,
                        detector_image_publisher_.publisher_yolo_left);
                };
                image_subscriber_.subscriber_left_image->AddReceiveCallback(
                    lambda_left);

                auto lambda_right =
                    [&](const char*, const xcmg_proto::OpencvImage& data_right,
                        long long, long long, long long) {
                        CallbackImageDetect(
                            params_.right_image_topic.c_str(), data_right, 0, 0,
                            0, detector_image_publisher_.publisher_yolo_right);
                    };
                image_subscriber_.subscriber_right_image->AddReceiveCallback(
                    lambda_right);

                return true;
            }

            void StreamProcessorStage::CallbackImageDetect(
                const char* topic_name_, const xcmg_proto::OpencvImage& data,
                long long time_, long long clock_, long long id,
                std::shared_ptr<
                    eCAL::protobuf::CPublisher<xcmg_proto::OpencvImage>>
                    pub) {
                mt_.lock();
                if (data.rows() > 0 && data.cols() > 0) {
                    ADEBUG_F("Received image of topic is {}", topic_name_);
                    cv::Mat image(data.rows(), data.cols(), data.elt_type());
                    cv::Mat result_img;
                    image.data = (uchar*)data.mat_data().data();

                    if (params_.model_flag) {
                        detector_ptr_->DetectCamera(image, result_img);
                        if (params_.ignore_region_flag) {
                            if (params_.camera_masks.find(topic_name_) ==
                                params_.camera_masks.end()) {
                                AWARN_F("No mask found for camera: {}",
                                        topic_name_);
                            } else {
                                if (params_.camera_masks[topic_name_].size() ==
                                    image.size()) {
                                    image.copyTo(
                                        result_img,
                                        params_.camera_masks[topic_name_]);
                                } else {
                                    AWARN_F(
                                        "Mask size mismatch: mask({},{}), "
                                        "image({},{})",
                                        params_.camera_masks[topic_name_].rows,
                                        params_.camera_masks[topic_name_].cols,
                                        image.rows, image.cols);
                                }
                            }
                        }
                        if (params_.mask_show_flag) {
                            cv::imshow(std::string(topic_name_) + "_mask",
                                       params_.camera_masks[topic_name_]);
                            cv::imshow(std::string(topic_name_) + "_origin",
                                       image);
                            cv::imshow(std::string(topic_name_) + "_result_img",
                                       result_img);
                            cv::waitKey(1);
                        }
                        pub->Send(data_converter_->ConvertImageMsg(result_img));
                    } else {
                        AERROR_F("Model is not loaded");
                    }

                    if (params_.rtsp_enable) {
                        PushRtsp(topic_name_, image);
                    }
                } else {
                    AERROR_F("Error image data");
                }
                mt_.unlock();
            }

            void StreamProcessorStage::PushRtsp(std::string topic_name,
                                                cv::Mat img) {
                if (img.rows != params_.rtsp_height ||
                    img.cols != params_.rtsp_width) {
                    cv::resize(
                        img, img,
                        cv::Size(params_.rtsp_width, params_.rtsp_height), 0, 0,
                        cv::INTER_LINEAR);
                }

                if (strcmp(topic_name.c_str(),
                           params_.front_image_topic.c_str()) == 0) {
                    rtsp_publisher_.writer_front.write(img);
                    ADEBUG_F("Push rtsp image of topic is {}", topic_name);
                } else if (strcmp(topic_name.c_str(),
                                  params_.back_image_topic.c_str()) == 0) {
                    rtsp_publisher_.writer_back.write(img);
                    ADEBUG_F("Push rtsp image of topic is {}", topic_name);
                } else if (strcmp(topic_name.c_str(),
                                  params_.left_image_topic.c_str()) == 0) {
                    rtsp_publisher_.writer_left.write(img);
                    ADEBUG_F("Push rtsp image of topic is {}", topic_name);
                } else if (strcmp(topic_name.c_str(),
                                  params_.right_image_topic.c_str()) == 0) {
                    rtsp_publisher_.writer_right.write(img);
                    ADEBUG_F("Push rtsp image of topic is {}", topic_name);
                } else {
                    AERROR_F("Error rtsp image topic name");
                }
            }
        }   // namespace stage
    }       // namespace core
}   // namespace xict_cpp
